#!/bin/bash

set -e

ECS_IMAGE=$1
VNC_PASSWORD=${@: -1}

sed -i 's/^#Port 22/Port 5000/' /etc/ssh/sshd_config
sed -i 's/^PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config

systemctl restart sshd

initialize(){
    expect <<EOF
        set timeout -1

        spawn /sbin/runuser -l root -c '/usr/bin/vncserver'

        expect {
            "Password:" { send "${VNC_PASSWORD}\n";exp_continue }
            "Verify:" { send "${VNC_PASSWORD}\n";exp_continue }
            "(y/n)?" { send 'n\n' }
        }

        expect eof
EOF
    /sbin/runuser -l root -c '/usr/bin/vncserver -kill :1'

    sed -i '$astartxfce4 &' /root/.vnc/xstartup && chmod +x /root/.vnc/xstartup && /sbin/runuser -l root -c '/usr/bin/vncserver :1'
}

case ${ECS_IMAGE} in
    *Ubuntu*|*Debian*)
        apt update -y && apt install -y xfce4 lightdm-
        apt install -y xfce4-goodies tightvncserver expect && initialize ${VNC_PASSWORD}

        cat > /etc/init.d/vncserver <<EOF
#! /bin/sh

### BEGIN INIT INFO
# Provides: vncserver
# Required-Start: $remote_fs $syslog
# Required-Stop: $remote_fs $syslog
# Default-Start: 2 3 4 5
# Default-Stop:
# Short-Description: vncserver
### END INIT INFO

set -e

case "\$1" in
    start)
        /sbin/runuser -l root -c '/usr/bin/vncserver :1'
        ;;
    stop)
        /sbin/runuser -l root -c '/usr/bin/vncserver -kill :1'
        ;;
    *)
        echo 'Usage: /etc/init.d/vncserver {start|stop}'
        exit 1
esac

exit 0
EOF
        chmod +x /etc/init.d/vncserver && update-rc.d vncserver defaults && systemctl enable vncserver
        ;;
    *CentOS*|*AlmaLinux*|*Rocky*)
        yum groupinstall -y xfce && yum install -y tigervnc-server expect && initialize ${VNC_PASSWORD}

        sed -i '$a:1=root' /etc/tigervnc/vncserver.users && systemctl enable vncserver@:1.service
        ;;
    *)
        echo "Unsupported image type!!!"
    exit 1
esac

exit 0
